﻿using System;
using System.Collections.Generic;

namespace ScapeGoatTree
{
    public class ScapeGoatTreeBasedSet
    {
        private readonly Stack<ScapeGoatNode> parents = new Stack<ScapeGoatNode>();

        private ScapeGoatNode root = null;

        public int Size { get; private set; }

        public bool Contains(int element)
        {
            var node = root;
            while (node != null)
            {
                if (element > node.Element)
                    node = node.Right;
                else if (element < node.Element)
                    node = node.Left;
                else
                    return true;
            }
            return false;
        }

        public void Add(int element)
        {
            parents.Clear();
            if (!AddElement(ref root, element))
                return;
            RefreshSize();
            if (!IsDeep())
                return;
            var child = parents.Pop();
            var childSize = 1;
            var scapeGoat = child;
            var scapeGoatSize = childSize;
            while (parents.Count != 0)
            {
                var parent = parents.Pop();
                var anotherChildSize = SizeOf(AnotherChild(parent, child));
                var parentSize = childSize + anotherChildSize + 1;
                if (!IsWeightBalanced(parentSize, childSize, anotherChildSize))
                {
                    scapeGoat = parent;
                    scapeGoatSize = parentSize;
                    break;
                }
                child = parent;
                childSize = parentSize;
            }
            var head = Flatten(scapeGoat, null);
            var scapeGoatParent = parents.Count == 0 ? null : parents.Pop();
            if (scapeGoatParent == null)
                root = BuildTree(ref head, 0, scapeGoatSize - 1);
            else if (scapeGoatParent.Right == scapeGoat)
                scapeGoatParent.Right = BuildTree(ref head, 0, scapeGoatSize - 1);
            else if (scapeGoatParent.Left == scapeGoat)
                scapeGoatParent.Left = BuildTree(ref head, 0, scapeGoatSize - 1);
            else
                throw new Exception(":(");
        }

        private static bool IsWeightBalanced(int parentSize, int childSize, int anotherChildSize)
        {
            var magicNumber = parentSize*0.5;
            if (childSize > magicNumber)
                return false;
            if (anotherChildSize > magicNumber)
                return false;
            return true;
        }

        private static ScapeGoatNode AnotherChild(ScapeGoatNode parent, ScapeGoatNode child)
        {
            return parent.Left == child ? parent.Right : parent.Left;
        }

        private static int SizeOf(ScapeGoatNode node)
        {
            if (node == null)
                return 0;
            return 1 + SizeOf(node.Left) + SizeOf(node.Right);
        }

        private bool IsDeep()
        {
            return parents.Count - 1 > (int) Math.Floor(Math.Log(Size, 2));
        }

        private void RefreshSize()
        {
            Size++;
        }

        private bool AddElement(ref ScapeGoatNode node, int element)
        {
            if (node == null)
            {
                node = new ScapeGoatNode(element);
                parents.Push(node);
                return true;
            }
            parents.Push(node);
            if (element > node.Element)
                return AddElement(ref node.Right, element);
            if (element < node.Element)
                return AddElement(ref node.Left, element);
            return false;
        }

        private static ScapeGoatNode Flatten(ScapeGoatNode node, ScapeGoatNode head)
        {
            if (node == null)
            {
                if (head != null)
                    head.Left = null;
                return head;
            }
            node.Right = Flatten(node.Right, head);
            var result = Flatten(node.Left, node);
            result.Left = null;
            return result;
        }

        private static ScapeGoatNode BuildTree(ref ScapeGoatNode head, int start, int end)
        {
            if (start > end)
                return null;
            var mid = start + (end - start)/2;
            var leftChild = BuildTree(ref head, start, mid - 1);
            var parent = head;
            parent.Left = leftChild;
            head = head.Right;
            parent.Right = BuildTree(ref head, mid + 1, end);
            return parent;
        }

        #region Nested type: ScapeGoatNode

        internal class ScapeGoatNode
        {
            public readonly int Element;
            public ScapeGoatNode Left;
            public ScapeGoatNode Right;

            public ScapeGoatNode(int element)
            {
                Element = element;
                Left = Right = null;
            }
        }

        #endregion
    }
}